import numpy as np


class DataLoader:

    def __init__(self, dataset, shuffle, collate_fn, batch_size,test=False):
        self.dataset = dataset
        self.shuffle = shuffle
        self.collate_fn = collate_fn
        self.batch_size = batch_size
        self.is_test=test

    def __next__(self):
        indics = np.arange(len(self.dataset))
        if self.shuffle:
            np.random.shuffle(indics)
        batch = []
        for i in range(0, len(self.dataset), self.batch_size):
            record = []
            for j in range(i, min(i + self.batch_size, len(self.dataset))):
                _index = indics[j]
                if self.is_test:
                    x1, x2 = self.dataset[_index]
                    record.append([x1, x2])                    
                else:
                    x1, x2, label = self.dataset[_index]
                    record.append([x1, x2, label])
            batch.append(self.collate_fn(record))
        return batch
